import torch
import torch.nn as nn
from nets import network_normalized


class Tandem(nn.Module):
    def __init__(self, inn_size, fnn_size, learning_rate=0.0001):
        super(Tandem, self).__init__()
        self.fnn_size = fnn_size
        self.inn_size = inn_size
        self.lr = learning_rate
        self.inn = network_normalized.Inv(inn_size=self.inn_size, learning_rate=self.lr)
        self.fnn = network_normalized.Fwd(fnn_size=self.fnn_size, learning_rate=self.lr)

    def forward(self, x):
        x = self.inn(x)
        x = self.fnn(x)
        return x

    def train(self, mode=True):
        self.fnn.eval()
        self.inn.train(mode)
        return self

    def eval(self):
        self.fnn.eval()
        self.inn.eval()
        return self

    def load_fnn(self, fnn_path: str):
        self.fnn.load_state_dict(torch.load(fnn_path))
        return self
